import pandas as pd
import json
from datetime import datetime
from utils.db import get_mysql_conn, get_redis_conn, get_kafka_consumer
import time
import threading

class DataCollector:
    """数据采集类，用于从Kafka消息队列收集用户行为数据"""
    
    def __init__(self):
        self.consumer = get_kafka_consumer()
        self.mysql_conn = get_mysql_conn()
        self.redis_conn = get_redis_conn()
        self.should_run = True
        
    def start_collection(self):
        """启动数据收集"""
        if not self.consumer:
            print("Kafka消费者未初始化，无法开始数据收集")
            return False
            
        print("开始数据收集...")
        
        # 创建一个线程来运行数据收集，以便不阻塞主线程
        self.collection_thread = threading.Thread(target=self._collect_data)
        self.collection_thread.daemon = True
        self.collection_thread.start()
        return True
        
    def stop_collection(self):
        """停止数据收集"""
        self.should_run = False
        if self.collection_thread:
            self.collection_thread.join()
        if self.consumer:
            self.consumer.close()
        print("数据收集已停止")
        
    def _collect_data(self):
        """从Kafka收集数据并存入数据库"""
        try:
            for message in self.consumer:
                if not self.should_run:
                    break
                    
                # 获取消息内容
                data = message.value
                
                # 数据验证
                if not self._validate_data(data):
                    continue
                    
                # 保存用户行为数据到MySQL
                self._save_behavior_to_db(data)
                
                # 更新Redis中的实时数据
                self._update_realtime_data(data)
                
                print(f"处理消息: {data}")
        except Exception as e:
            print(f"数据收集过程中发生错误: {e}")
    
    def _validate_data(self, data):
        """验证数据有效性"""
        required_fields = ['user_id', 'event_type', 'timestamp', 'course_id']
        
        # 检查必要字段是否存在
        for field in required_fields:
            if field not in data:
                print(f"数据缺少必要字段: {field}")
                return False
                
        # 验证时间戳有效性（不能是未来时间）
        try:
            event_time = datetime.fromtimestamp(data['timestamp'])
            current_time = datetime.now()
            if event_time > current_time:
                print(f"无效的未来时间戳: {event_time}")
                return False
        except:
            print(f"无效的时间戳格式: {data['timestamp']}")
            return False
            
        # 验证事件类型有效性
        valid_events = ['view', 'click', 'favorite', 'rate', 'complete', 'purchase']
        if data['event_type'] not in valid_events:
            print(f"无效的事件类型: {data['event_type']}")
            return False
            
        return True
        
    def _save_behavior_to_db(self, data):
        """保存用户行为数据到MySQL数据库"""
        try:
            if not self.mysql_conn:
                self.mysql_conn = get_mysql_conn()
                
            with self.mysql_conn.cursor() as cursor:
                sql = """
                INSERT INTO user_behaviors (user_id, course_id, event_type, timestamp, extra_data)
                VALUES (%s, %s, %s, %s, %s)
                """
                
                # 将额外数据转为JSON
                extra_data = {}
                for key, value in data.items():
                    if key not in ['user_id', 'course_id', 'event_type', 'timestamp']:
                        extra_data[key] = value
                        
                cursor.execute(
                    sql, 
                    (
                        data['user_id'],
                        data['course_id'],
                        data['event_type'],
                        datetime.fromtimestamp(data['timestamp']),
                        json.dumps(extra_data)
                    )
                )
                self.mysql_conn.commit()
        except Exception as e:
            print(f"保存用户行为数据失败: {e}")
            if self.mysql_conn:
                self.mysql_conn.rollback()
    
    def _update_realtime_data(self, data):
        """更新Redis中的实时数据"""
        try:
            if not self.redis_conn:
                self.redis_conn = get_redis_conn()
                
            # 更新课程热度计数
            if data['event_type'] == 'view':
                self.redis_conn.hincrby(f"course:{data['course_id']}:counts", "views", 1)
            elif data['event_type'] == 'purchase':
                self.redis_conn.hincrby(f"course:{data['course_id']}:counts", "purchases", 1)
            elif data['event_type'] == 'favorite':
                self.redis_conn.hincrby(f"course:{data['course_id']}:counts", "favorites", 1)
                
            # 更新用户最近行为
            user_recent_key = f"user:{data['user_id']}:recent_behaviors"
            self.redis_conn.lpush(user_recent_key, json.dumps(data))
            self.redis_conn.ltrim(user_recent_key, 0, 99)  # 只保留最近100条行为
            
        except Exception as e:
            print(f"更新Redis实时数据失败: {e}")

# 测试代码
if __name__ == "__main__":
    collector = DataCollector()
    if collector.start_collection():
        try:
            # 保持程序运行
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            collector.stop_collection()
            print("程序已终止") 